/*
 * Copyright 2012 Niklas Rehfeld.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.madpie.transport;

import java.io.File;
import java.io.FileNotFoundException;
import java.util.Scanner;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 *
 * @author Niklas Rehfeld
 */
public class SshTransfer
{

    public static final String CPU_COMMAND = "cpu";
    public static final String HW_COMMAND = "hw";
    public static final String DISK_COMMAND = "disk";
    public static final String DISK_IMAGE_COMMAND = "image";
    public static final String TEST_CONNECTION_COMMAND = "test";
    public static final String DISK_CHECKSUM_COMMAND = "checksum";
    private String userName, targetHost;
    //TODO get these from somewhere. at the moment we have these as consts in two classes.
    public static final String CPU_FILE = "cpu.info";
    public static final String HW_FILE = "hw.info";
    public static final String DISK_FILE = "disks.info";
    private String baseDir;

    /**
     *
     * @param userName
     * @param targetHost
     * @param outputDir
     */
    public SshTransfer(String userName, String targetHost, String outputDir)
    {
        this.userName = userName;
        this.targetHost = targetHost;
        this.baseDir = outputDir;
    }

    /**
     * Main method.
     *
     * TODO : make a config file that allows the user to set the commands.
     *
     * @param args [username host outputDir command]
     */
    public static void main(String[] args)
    {
        Logger.getLogger(SshTransfer.class.getName()).setLevel(Level.ALL);

        if (args.length < 4)
        {
            Logger.getLogger(SshTransfer.class.getName()).log(Level.SEVERE,
                    "Not enough args. usage: \n\t "
                    + "sshtransfer username host outputDir command [command_opts]");
            System.exit(-1);
        }
        SshTransfer trans = new SshTransfer(args[0], args[1], args[2]);
        String command = args[3];
        boolean result = false;
        if (command.equalsIgnoreCase(CPU_COMMAND))
        {
            result = trans.getCPUInfo();
        }
        else if (command.equalsIgnoreCase(HW_COMMAND))
        {
            result = trans.getHwInfo();
        }
        else if (command.equalsIgnoreCase(DISK_COMMAND))
        {
            result = trans.getDiskInfoFdisk();
        }
        else if (command.equalsIgnoreCase(DISK_IMAGE_COMMAND))
        {
            if (args.length < 5)
            {
                Logger.getLogger(SshTransfer.class.getName()).log(Level.SEVERE,
                        "Must specify a disk to image.");
                System.exit(-2);
            }
            result = trans.getDiskImage(args[4]);
        }
        else if (command.equalsIgnoreCase(TEST_CONNECTION_COMMAND))
        {
            result = trans.testConnection();
        }
        else if (command.equalsIgnoreCase(DISK_CHECKSUM_COMMAND))
        {
            if (args.length < 5)
            {
                Logger.getLogger(SshTransfer.class.getName()).log(Level.SEVERE,
                        "Must specify a disk to image.");
                System.exit(-2);
            }
            result = trans.getDiskChecksum(args[4]);
        }
        else
        {
            Logger.getLogger(SshTransfer.class.getName()).log(Level.SEVERE,
                    "Unrecognised command.");
            System.exit(100);
        }
        System.exit(result ? 0 : -1);

    }

    private boolean testConnection()
    {
        SshTermWrapper t = new SshTermWrapper(userName, targetHost);
        return t.testSshConnection();
    }

    /**
     * Retrieves the hardware information from the server via SSH.
     *
     * TODO : make it actually test the output.
     *
     * @return true if it succeeded.
     *
     */
    public boolean getHwInfo()
    {
        File out = new File(baseDir, HW_FILE);
        SshTermWrapper s = new SshTermWrapper(userName, targetHost);
        Process p = s.
                runSshInTerminal(
                "'lspci -vvmmnn' > "
                + out.getAbsolutePath());
        try
        {
            p.waitFor();
            if (p.exitValue() != 0)
            {
                Logger.getLogger(SshTransfer.class.getName()).log(Level.WARNING,
                        "Command failed.");
                return false;
            }

        } catch (InterruptedException ex)
        {
            Logger.getLogger(SshTransfer.class.getName()).log(Level.SEVERE,
                    "Interrupted while waiting for hardware info", ex);
        }
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "created {0}", out.getPath());
        return (p.exitValue() == 0);

    }

    /**
     * Retrieves the cpu information from the server via SSH.
     *
     * @return true if the command succeeded.
     */
    public boolean getCPUInfo()
    {
        SshTermWrapper s = new SshTermWrapper(userName, targetHost);
        File out = new File(baseDir, CPU_FILE);
        Process p = s.
                runSshInTerminal(
                "' lscpu ' > "
                + out.getAbsolutePath());
        try
        {
            p.waitFor();

        } catch (InterruptedException ex)
        {
            Logger.getLogger(SshTransfer.class.getName()).log(Level.SEVERE,
                    "Interrupted while waiting for metadata", ex);
            return false;
        }
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "created {0}", out.getPath());
        return true;
    }

    /**
     * Uses fdisk to get file system information off a remote computer. user
     * needs to have sufficient permissions on the remote computer.
     *
     * @return
     */
    public boolean getDiskInfoFdisk()
    {
        SshTermWrapper s = new SshTermWrapper(userName, targetHost);
        File out = new File(baseDir, DISK_FILE);
        Process p = s.
                runSshInTerminal(
                "'fdisk -l' > "
                + out.getAbsolutePath());
        try
        {
            p.waitFor();

        } catch (InterruptedException ex)
        {
            Logger.getLogger(SshTransfer.class.getName()).log(Level.SEVERE,
                    "Interrupted while waiting for metadata", ex);
            return false;
        }
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "created {0}", out.getPath());
        return true;
    }

    /**
     * runs
     * <code>dd <em>disk_path</em></code> and pipes the output into a file
     * called <em>disk_path.img</em>
     *
     * Note that unlike the other ones, this one returns immediately, without
     * waiting for the transfer to complete. The amount done can be determined
     * by looking at the size of the output file and comparing it to the size
     * that was reported by fdisk earlier in the workflow.
     *
     * The user on the remote computer must have sufficient permissions.
     *
     * @param disk the path to the disk e.g. /dev/sda.
     * @return true, unless something really unexpected happens...
     */
    private boolean getDiskImage(String disk)
    {
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "Transferring disk ''{0}''", disk);
        SshTermWrapper s = new SshTermWrapper(userName, targetHost);
        String filename = disk.replaceAll("/", "_");
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "Filename: " + filename);
        File out = new File(baseDir, filename + ".img");
        Process p = s.
                runSshInTerminal(
                "'dd if=" + disk + "' > "
                + out.getAbsolutePath());
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "created {0}", out.getPath());
        return true;
    }

    private boolean getDiskChecksum(String disk)
    {
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "Transferring disk ''{0}''", disk);
        SshTermWrapper s = new SshTermWrapper(userName, targetHost);
        String filename = disk.replaceAll("/", "_");
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "Filename: " + filename);
        File out = new File(baseDir, filename + ".md5");
        //run md5sum, redirect transfer stats to _dev_bar.md5.err and output sum 
        //to _dev_bar.md5
        Process p = s.
                runSshInTerminal(
                "'dd if=" + disk + " | md5sum' 2&> "
                + out.getAbsolutePath() + ".err >" + out.getAbsolutePath());
        Logger.getLogger(SshTransfer.class.getName()).log(Level.INFO,
                "created {0}", out.getPath());
        try
        {
            p.waitFor();
        } catch (InterruptedException ex)
        {
            Logger.getLogger(SshTransfer.class.getName()).
                    log(Level.SEVERE, "Interrupted!", ex);
        }
        Scanner scan;
        try
        {
            scan = new Scanner(new File(out.getAbsolutePath() + ".err"));
            while (scan.hasNextLine())
            {
                if (scan.nextLine().matches("^dd: reading .* error$"))
                {
                    return false;
                }
            }

        } catch (FileNotFoundException ex)
        {
            Logger.getLogger(SshTransfer.class.getName()).
                    log(Level.SEVERE, null, ex);
        }
        return true;
    }
}
